{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "722be46f-0441-4675-a79e-856313859ec4",
   "metadata": {},
   "source": [
    "# Fine-Tune LLM with Data Cache\n",
    "\n",
    "This notebook demonstrates how to fine-tune a BERT model with Kubeflow TrainJob using data streaming from a distributed cache.\n",
    "\n",
    "We will use an iceberg table with [the Amazon reviews dataset](https://huggingface.co/datasets/fancyzhx/amazon_polarity) loaded from data cache cluster, which provides efficient data streaming for distributed training workloads."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98295561-e413-4b09-8393-3a0a11b57b63",
   "metadata": {},
   "source": [
    "## Install Kubeflow SDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39a7b539-0d4b-44f6-aa96-f42ff4cb846b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the Kubeflow SDK\n",
    "!pip install -U kubeflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bef527f-121a-4c8b-8c99-2a63b299f638",
   "metadata": {},
   "source": [
    "## List Available Training Runtimes\n",
    "\n",
    "Get available Kubeflow Training Runtimes using the list_runtimes() API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b953ecca-b22e-42d8-8fa1-2168fe7a6e8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kubeflow.trainer import TrainerClient\n",
    "\n",
    "client = TrainerClient()\n",
    "for runtime in client.list_runtimes():\n",
    "    print(f'{runtime}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec94f65c-1cc6-47ff-ad8d-92c35776f775",
   "metadata": {},
   "source": [
    "## Define Training Function with DataCacheDataset\n",
    "\n",
    "This training function uses DataCacheDataset to stream data from the cache cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96fd1dc0-63ea-4eb1-919a-14257eca20d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_func():\n",
    "    import time\n",
    "    import datetime\n",
    "    from multiprocessing import current_process\n",
    "    import torch.distributed as dist\n",
    "    import pyarrow as pa\n",
    "    import pyarrow.flight\n",
    "    from pyarrow._flight import FlightClient\n",
    "    import torch\n",
    "    import os\n",
    "    from torch.utils.data import IterableDataset\n",
    "    from transformers import DistilBertForSequenceClassification, DistilBertTokenizer, AdamW\n",
    "\n",
    "    \n",
    "    # Define DataCacheDataset to access data from cache\n",
    "    class DataCacheDataset(IterableDataset):\n",
    "        def __init__(self, seed=0, batch_size=128):\n",
    "            train_job_name = os.getenv(\"TRAIN_JOB_NAME\")\n",
    "            self.endpoint = f'grpc://{train_job_name}-cache-service:50051'\n",
    "            self.seed = seed\n",
    "            self.batch_size = batch_size\n",
    "\n",
    "        def set_epoch(self, epoch):\n",
    "            self.seed = self.seed + epoch\n",
    "\n",
    "        def from_arrow_rb_to_tensor(self, chunk):\n",
    "            return\n",
    "\n",
    "        # Extract RecordBatches and split them into chunks\n",
    "        def stream_recordbatch_chunks(self, recordbatch_stream, chunk_size):\n",
    "            leftover = None\n",
    "        \n",
    "            for rb in recordbatch_stream:\n",
    "                batch = rb.data\n",
    "                if leftover is not None:\n",
    "                    arrays = [\n",
    "                        pa.concat_arrays([leftover.column(i), batch.column(i)])\n",
    "                        for i in range(len(batch.columns))\n",
    "                    ]\n",
    "                    batch = pa.RecordBatch.from_arrays(arrays, batch.schema.names)\n",
    "\n",
    "                for i in range(0, batch.num_rows, chunk_size):\n",
    "                    chunk = batch.slice(i, chunk_size)\n",
    "        \n",
    "                    if chunk.num_rows < chunk_size:\n",
    "                        leftover = chunk\n",
    "                        break\n",
    "                    else:\n",
    "                        yield from self.from_arrow_rb_to_tensor(chunk)\n",
    "                else:\n",
    "                    leftover = None\n",
    "        \n",
    "            if leftover is not None and leftover.num_rows > 0:\n",
    "                yield from self.from_arrow_rb_to_tensor(leftover)\n",
    "\n",
    "        def __iter__(self):\n",
    "            worker_info = torch.utils.data.get_worker_info()\n",
    "            \n",
    "            if worker_info is None:\n",
    "                worker_id = 0\n",
    "                num_workers = 1\n",
    "            else:\n",
    "                worker_id = worker_info.id\n",
    "                num_workers = worker_info.num_workers\n",
    "\n",
    "            world_size = dist.get_world_size()\n",
    "            rank = dist.get_rank()\n",
    "                \n",
    "            # Calculate shard index\n",
    "            total_workers = num_workers * world_size\n",
    "            global_worker_id = worker_id * world_size + rank\n",
    "            index = (global_worker_id + self.seed) % total_workers\n",
    "\n",
    "            descriptor = pa.flight.FlightDescriptor.for_path(*[s for s in [str(index), str(total_workers)]])\n",
    "            client = pa.flight.connect(self.endpoint)\n",
    "            flight_info = client.get_flight_info(descriptor)\n",
    "\n",
    "            for endpoint in flight_info.endpoints:\n",
    "                for location in endpoint.locations:\n",
    "                    _client = FlightClient(location.uri)\n",
    "                    ticket = endpoint.ticket\n",
    "                    reader = _client.do_get(ticket)\n",
    "                    \n",
    "                    yield from self.stream_recordbatch_chunks(reader, self.batch_size)\n",
    "\n",
    "    # Define dataset for Amazon reviews\n",
    "    class AmazonReviewDataset(DataCacheDataset):\n",
    "        def __init__(self, tokenizer, seed=0):\n",
    "            super().__init__(seed)\n",
    "            self.tokenizer = tokenizer\n",
    "            self.row_count = 0\n",
    "\n",
    "        def from_arrow_rb_to_tensor(self, chunk):\n",
    "            self.row_count = self.row_count + chunk.num_rows\n",
    "            texts = chunk.column(\"content\").to_pylist()\n",
    "            labels = chunk.column(\"label\").to_pylist()\n",
    "\n",
    "            encoding = self.tokenizer(\n",
    "                texts,\n",
    "                padding=\"max_length\",\n",
    "                truncation=True,\n",
    "                max_length=128,\n",
    "                return_tensors=\"pt\"\n",
    "            )\n",
    "        \n",
    "            yield {\n",
    "                \"input_ids\": encoding[\"input_ids\"].squeeze(0),\n",
    "                \"attention_mask\": encoding[\"attention_mask\"].squeeze(0),\n",
    "                \"labels\": torch.tensor(labels)\n",
    "            }\n",
    "\n",
    "    device, backend = (\"cuda\", \"nccl\") if torch.cuda.is_available() else (\"cpu\", \"gloo\")\n",
    "    print(f\"Using Device: {device}, Backend: {backend}\")\n",
    "\n",
    "    local_rank = int(os.getenv(\"LOCAL_RANK\", 0))\n",
    "    dist.init_process_group(backend=backend)\n",
    "    print(\n",
    "        \"Distributed Training for WORLD_SIZE: {}, RANK: {}, LOCAL_RANK: {}\".format(\n",
    "            dist.get_world_size(),\n",
    "            dist.get_rank(),\n",
    "            local_rank,\n",
    "        )\n",
    "    )\n",
    "    device = torch.device(f\"{device}:{local_rank}\")\n",
    "    model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=5)\n",
    "    model.to(device)\n",
    "    optimizer = AdamW(model.parameters(), lr=2e-5)\n",
    "\n",
    "    tokenizer = DistilBertTokenizer.from_pretrained(\"distilbert-base-uncased\")\n",
    "\n",
    "    start = time.time()\n",
    "    dataset = AmazonReviewDataset(tokenizer)\n",
    "    dataloader = torch.utils.data.DataLoader(dataset, batch_size=None)\n",
    "\n",
    "    # Training loop\n",
    "    for epoch in range(1):\n",
    "        dataset.set_epoch(epoch)\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            attention_mask = batch[\"attention_mask\"].to(device)\n",
    "            labels = batch[\"labels\"].to(device)\n",
    "            outputs = model(input_ids, attention_mask=attention_mask, labels=labels)\n",
    "            loss = outputs.loss\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            if dist.get_rank() == 0:\n",
    "                print(f\"Epoch {epoch}, Loss: {loss.item()}\")\n",
    "\n",
    "    dist.barrier()\n",
    "    if dist.get_rank() == 0:\n",
    "        print(\"Training is finished\")\n",
    "\n",
    "    end = time.time()\n",
    "    print(f\"Time taken: {end - start:.6f} seconds\")\n",
    "    print(f\"Total rows processed: {dataset.row_count}\")\n",
    "\n",
    "    dist.destroy_process_group()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bda5814-252c-418c-ac1d-297b8953e747",
   "metadata": {},
   "source": [
    "## Create TrainJob with Data Cache\n",
    "\n",
    "Use the train() API with DataCacheInitializer to create a TrainJob that streams data from cache."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d76db904-d597-4e7e-b37b-9ec1762d2482",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kubeflow.trainer import CustomTrainer, DataCacheInitializer, Initializer\n",
    "\n",
    "# Get the cache runtime\n",
    "torch_cache_runtime = None\n",
    "for runtime in client.list_runtimes():\n",
    "    if runtime.name == \"torch-distributed-with-cache\":\n",
    "        torch_cache_runtime = runtime\n",
    "        break\n",
    "\n",
    "if torch_cache_runtime is None:\n",
    "    raise ValueError(\"torch-distributed-with-cache runtime not found. Please ensure the runtime is installed.\")\n",
    "\n",
    "# Configure your data cache settings\n",
    "# Replace these values with your actual metadata location and storage URI\n",
    "METADATA_LOC = \"Iceberg table metadata file path on s3\"\n",
    "STORAGE_URI = \"cache://<schema_name>/<table_name>\"\n",
    "IAM_ROLE= \"IAM role with access to table\"\n",
    "\n",
    "job_name = client.train(\n",
    "    trainer=CustomTrainer(\n",
    "        func=train_func,\n",
    "        # Set how many PyTorch nodes you want to use for distributed training\n",
    "        num_nodes=8,\n",
    "        # Set the resources for each PyTorch node\n",
    "        resources_per_node={\n",
    "            \"cpu\": 3,\n",
    "            \"memory\": \"8Gi\",\n",
    "            # Uncomment to use GPUs\n",
    "            \"nvidia.com/gpu\": 1,\n",
    "        },\n",
    "        packages_to_install=[\"pyarrow==19.0.0\", \"transformers==4.45.2\"]\n",
    "    ),\n",
    "    # Dataset config with data cache initializer\n",
    "    initializer=Initializer(\n",
    "        dataset=DataCacheInitializer(\n",
    "            num_data_nodes=4,\n",
    "            metadata_loc=METADATA_LOC,\n",
    "            storage_uri=STORAGE_URI,\n",
    "            iam_role=IAM_ROLE\n",
    "        )\n",
    "    ),\n",
    "    runtime=torch_cache_runtime\n",
    ")\n",
    "\n",
    "print(f\"TrainJob created: {job_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c531106f-f833-45ad-abe2-28858b699bfc",
   "metadata": {},
   "source": [
    "## Check the TrainJob Info\n",
    "\n",
    "Use the list_jobs() and get_job() APIs to get information about the created TrainJob."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0091971d-f370-4a2f-9986-8ccd2a75bed9",
   "metadata": {},
   "outputs": [],
   "source": [
    "for job in client.list_jobs():\n",
    "    print(f\"TrainJob: {job.name}, Status: {job.status}, Created at: {job.creation_timestamp}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe579beb-c97b-496d-bc07-ba9dbac9562a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get details about the TrainJob steps\n",
    "for step in client.get_job(name=job_name).steps:\n",
    "    print(f\"Step: {step.name}, Status: {step.status}, Devices: {step.device} x {step.device_count}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f86e51f7-eb3e-40c4-9756-1bfbb304368c",
   "metadata": {},
   "source": [
    "## Watch the TrainJob Logs\n",
    "\n",
    "Use the get_job_logs() API to follow the TrainJob logs in real-time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b176b43e-3276-4568-9333-915e0a69bac0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wait for the job to reach running state\n",
    "client.wait_for_job_status(job_name, status={\"Running\"})\n",
    "print(\"Job is in running state\")\n",
    "\n",
    "# Follow the logs\n",
    "for line in client.get_job_logs(job_name, follow=True):\n",
    "    print(line)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d29925f6-2e6b-4db1-80b7-49fedd5910e9",
   "metadata": {},
   "source": [
    "## Delete the TrainJob\n",
    "\n",
    "When the TrainJob is finished, you can delete the resource."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6db3f918-0fce-447f-8b21-dbcaf75492dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to delete the job\n",
    "# client.delete_job(job_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0941b042-da4b-4551-8f8c-5e8dc792d850",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
